# 导入所需的库
from langchain.prompts import PromptTemplate
from langchain.chat_models import ChatOpenAI
from langchain.agents import initialize_agent, Tool
from langchain.agents import AgentType
from langchain_ollama import OllamaLLM
import yaml
from langchain_community.utilities import SerpAPIWrapper

if __name__ == "__main__":
    llm = OllamaLLM(model="qwen:latest")
    with open('config.yaml', 'r') as file:
        data = yaml.safe_load(file)

    # 寻找UID的模板
    template = """Given the flower "{flower}", I want you to find a related Weibo UID.
    Your answer must contain only the UID, which is a number.
    The UID should be part of a URL that starts with https://weibo.com/u/.
    For example, if the Weibo URL is https://weibo.com/u/1669879400, then the UID is 1669879400.
    This is just an example; please provide the actual UID related to the given flower.

    Important notes:
    - Do not return any text other than the UID.
    - Do not return the full URL, only the UID.
    - If you cannot find a specific UID related to the flower, return "Not Found".
    - Incorrect responses include: "I understand what you need.", "Here is the UID:", or any other text that is not a number.
    """
    # 完整的提示模板
    prompt_template = PromptTemplate(
        input_variables=["flower"], template=template
    )
    search = SerpAPIWrapper(serpapi_api_key=data['serpapi_api_key'])
    # 代理的工具
    tools = [
        Tool(
            name="Crawl Google for 微博 page",
            func=search.run,
            description="useful for when you need get the 微博 UID",
        )
    ]

    # 初始化代理
    agent = initialize_agent(
        tools, 
        llm, 
        agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, 
        verbose=True
    )


    # 返回找到的UID
    ID = agent.run(prompt_template.format_prompt(flower="玫瑰"))
